import os
from collections import namedtuple
from datetime import datetime

TrainRecord = namedtuple("TrainRecord", ["record_metric", "model_filename"])

class BestCheckpointSaver:
    def __init__(self, max_checkpoint_num=2, outputs="./outputs") -> None:
        self._outputs_dir = outputs
        self._max_checkpoint_num = max_checkpoint_num
        self._time_prefix = datetime.now().strftime(r"%m%d%H%M")
        
        if not os.path.isdir(self._outputs_dir):
            os.makedirs(self._outputs_dir)
        
        self._best_checkpoint = []
        pass

    def update(self, epoch, trainer, metric_value):
        
        # 命名规则：时间戳-性能指标
        target_filename = f"{self._time_prefix}-{metric_value:.4f}.ckpt"
        full_target_filepath = os.path.join(self._outputs_dir, target_filename)
        
        # 保存要求：没达到指定的保存数目，或者当前模型性能优于历史记录中的某条性能
        if len(self._best_checkpoint) < self._max_checkpoint_num:

            trainer.dump_checkpoint(epoch, metric_value, full_target_filepath)  # 保存当前模型
            self._best_checkpoint.append(TrainRecord(metric_value, full_target_filepath))# 当前模型添加到记录列表中

        elif metric_value > self._best_checkpoint[-1].record_metric:  # 当前模型的性能好于记录中最差的那个

                trainer.dump_checkpoint(epoch, metric_value,full_target_filepath)  # 保存当前的这个较好的模型
                try:
                    os.remove(self._best_checkpoint.pop().model_filename)  # 删除较差的那个模型
                except FileNotFoundError:
                    print("File Not Found.")
                    pass
                self._best_checkpoint.append(TrainRecord(metric_value, full_target_filepath))

        self._best_checkpoint = sorted(self._best_checkpoint, key=lambda x: x[0], reverse=True)  # 放在后面保证列表中至少有1条记录后排序
        return None
